package cn.whale.config;

import cn.whale.assistant.McpAssistant;
import cn.whale.assistant.OllamaAssistant;
import cn.whale.assistant.RAGAssistant;
import cn.whale.chain.TaskTypeAssistant;
import cn.whale.tools.WeatherTool;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.web.search.WebSearchEngine;
import dev.langchain4j.web.search.WebSearchTool;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class LLMConfig {


    /**
     * ollama大模型 - ai -service
     * 记忆功能
     */
    @Bean
    public OllamaAssistant ollamaAssistant(QwenStreamingChatModel streamingChatModel, WebSearchEngine webSearchEngine){

        //对话记忆功能实现
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder().chatMemoryStore(new PersistentChatMemoryStore()).maxMessages(10).build();

        //RAG检索
        return AiServices.builder(OllamaAssistant.class)
                //流式对话
                .streamingChatLanguageModel(streamingChatModel)
                //记忆功能
                .chatMemoryProvider((memoryId -> chatMemory))
                //调用自定义工具 ， web搜索工具
                .tools(new WeatherTool(),new WebSearchTool(webSearchEngine))
                .build();
    }
    /**
     * ollama大模型 - ai -service
     * 记忆功能
     */
    @Bean
    public RAGAssistant ragAssistant(QwenStreamingChatModel streamingChatModel, EmbeddingStore<TextSegment> embeddingStore){

        //对话记忆功能实现
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder().chatMemoryStore(new PersistentChatMemoryStore()).maxMessages(10).build();

        //内容检索
        EmbeddingStoreContentRetriever embeddingStoreContentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .maxResults(5)
                .minScore(0.7D)
                .build();

        //RAG检索
        return AiServices.builder(RAGAssistant.class)
                //RAG知识库检索
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                //流式对话
                .streamingChatLanguageModel(streamingChatModel)
                //记忆功能
                .chatMemoryProvider((memoryId -> chatMemory))
                .build();
    }

    @Bean
    public TaskTypeAssistant taskTypeAssistant(QwenChatModel qwenChatModel){
        return AiServices.builder(TaskTypeAssistant.class).chatLanguageModel(qwenChatModel).build();
    }

    @Bean
    public McpAssistant mcpAssistant(QwenStreamingChatModel qwenChatModel){
        return AiServices.builder(McpAssistant.class).streamingChatLanguageModel(qwenChatModel).build();
    }


}
